#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

import sys

default_encoding = 'utf-8'
if sys.getdefaultencoding() != default_encoding:
    reload(sys)
    sys.setdefaultencoding(default_encoding)

from pymongo import MongoClient
import re


def TestMatch(strr, regex):
    match = re.match(regex, strr)
    if match:
        return True
    else:
        return False


# 连接MongoDB
def InitMongoDB():
    MONGODB_HOST = 'localhost'
    MONGODB_PORT = 27017
    MONGODB_DB_NAME = 'pmc'
    MONGODB_COLLECTION_NAME = 'filelist'
    connection = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = connection[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return connection, collection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


if __name__ == "__main__":
    regex_with_oapackage = '^oa_package/[\da-f]{2}/[\da-f]{2}/PMC\d{5,8}\.tar\.gz$'
    regex_without_oapackage = '^[\da-f]{2}/[\da-f]{2}/PMC\d{5,8}\.tar\.gz$'

    MongoDB_Connection, MongoDB_Collection = InitMongoDB()

    for result in MongoDB_Collection.find({"type": "tarball"}):
        id = result["_id"]
        path = result["path"]
        pmcid = result["pmcid"]
        MongoDB_Collection.update_one({"_id": id}, {"$addToSet": {"paths": path}})
        if not TestMatch(path, regex_without_oapackage):
            path2 = path[:6] + "PMC" + str(pmcid) + ".tar.gz"
            MongoDB_Collection.update_one({"_id": id}, {"$addToSet": {"paths": path2}})

    CloseMongoDB(MongoDB_Connection)
